#!/usr/bin/python
from __future__ import print_function
import sys
import swsssdk
import redis
import argparse
from multiprocessing import Pool
from functools import partial

def handle_single_instance(op, use_unix_socket, inst_info):
    inst_hostname = inst_info['hostname']
    if use_unix_socket:
        inst_unix_socket_path = inst_info['unix_socket_path']
        r = redis.Redis(host=inst_hostname, unix_socket_path=inst_unix_socket_path)
    else:
        inst_port = inst_info['port']
        r = redis.Redis(host=inst_hostname, port=inst_port)
    rsp = False
    msg = 'Could not connect to Redis at {}:{}: Connection refused'.format(inst_hostname,
           inst_unix_socket_path if use_unix_socket else inst_port)
    try:
        if op == 'PING':
            rsp = r.ping()
        elif op == 'SAVE':
            rsp = r.save()
        elif op == 'FLUSHALL':
            rsp = r.flushall()
        else:
            raise ValueError("Operation {} is not supported".format(op))
    except redis.exceptions.ConnectionError as e:
        pass
    return None if rsp else msg

def handle_all_instances(namespace, op, use_unix_socket=False):
    # Use the tcp connectivity if namespace is local and unixsocket cmd_option is present.
    if namespace is not None:
        use_unix_socket = True

    db_insts = swsssdk.SonicDBConfig.get_instancelist(namespace)
    # Operate All Redis Instances in Parallel
    # TODO: if one of the operations failed, it could fail quickly and not necessary to wait all other operations
    p = Pool(len(db_insts))
    func = partial(handle_single_instance, op, use_unix_socket)
    rsps = p.map(func, [v for k, v in db_insts.items()])
    msg = [rsp for rsp in rsps if rsp]

    if msg:
        print('\n'.join(msg))
        sys.exit(1)

    if op == 'PING':
        print('PONG')
    else:
        print('OK')
    sys.exit(0)

def execute_cmd(dbname, cmd, namespace, use_unix_socket=False):
    if namespace is None:
        if use_unix_socket:
            dbconn = swsssdk.SonicV2Connector(use_unix_socket_path=True, decode_responses=True)
        else:
            dbconn = swsssdk.SonicV2Connector(use_unix_socket_path=False, decode_responses=True)
    else:
        dbconn = swsssdk.SonicV2Connector(use_unix_socket_path=True, namespace=namespace, decode_responses=True)
    try:
        dbconn.connect(dbname)
    except RuntimeError:
        msg = "Invalid database name input : '{}'".format(dbname)
        print(msg, file=sys.stderr)
        sys.exit(1)
    else:
        client = dbconn.get_redis_client(dbname)
        resp = client.execute_command(*cmd)
        """
        sonic-db-cli output format mimic the non-tty mode output format from redis-cli
        based on our usage in SONiC, None and list type output from python API needs to be modified
        with these changes, it is enough for us to mimic redis-cli in SONiC so far since no application uses tty mode redis-cli output
        """
        if resp is None:
            print()
        elif isinstance(resp, list):
            print("\n".join(resp))
        else:
            print(resp)
        sys.exit(0)

def main():
    parser = argparse.ArgumentParser(description='SONiC DB CLI',
                                     formatter_class=argparse.RawTextHelpFormatter,
                                     epilog=
"""
**sudo** needed for commands accesing a different namespace [-n], or using unixsocket connection [-s]

Example 1: sonic-db-cli -n asic0 CONFIG_DB keys \*
Example 2: sonic-db-cli -n asic2 APPL_DB HGETALL VLAN_TABLE:Vlan10
Example 3: sonic-db-cli APPL_DB HGET VLAN_TABLE:Vlan10 mtu
Example 4: sonic-db-cli -n asic3 APPL_DB EVAL "return {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}" 2 k1 k2 v1 v2
Example 5: sonic-db-cli PING | sonic-db-cli -s PING
Example 6: sonic-db-cli SAVE | sonic-db-cli -s SAVE
Example 7: sonic-db-cli FLUSHALL | sonic-db-cli -s FLUSHALL
""")
    parser.add_argument('-s', '--unixsocket', help="Override use of tcp_port and use unixsocket", action='store_true')
    parser.add_argument('-n', '--namespace', type=str, help="Namespace string to use asic0/asic1.../asicn", default=None)
    parser.add_argument('db_or_op', type=str, help='Database name Or Unary operation(only PING/SAVE/FLUSHALL supported)')
    parser.add_argument('cmd', nargs='*', type=str, help='Command to execute in database')

    args = parser.parse_args()

    if args.db_or_op:
        # Load the database config for the namespace
        if args.namespace is not None:
            swsssdk.SonicDBConfig.load_sonic_global_db_config(namespace=args.namespace)
        if args.cmd:
            execute_cmd(args.db_or_op, args.cmd, args.namespace, args.unixsocket)
        elif args.db_or_op in ('PING', 'SAVE', 'FLUSHALL'):
            # redis-cli doesn't depend on database_config.json which could raise some exceptions
            # sonic-db-cli catch all possible exceptions and handle it as a failure case which not return 'OK' or 'PONG'
            try:
                handle_all_instances(args.namespace, args.db_or_op, args.unixsocket)
            except Exception as ex:
                template = "An exception of type {0} occurred. Arguments:\n{1!r}"
                message = template.format(type(ex).__name__, ex.args)
                print(message, file=sys.stderr)
                sys.exit(1)
        else:
            parser.print_help()
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
